import psutil
import os
import time
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import argparse


def monitor_cpu_memory(process_id):
    process = psutil.Process(process_id)
    cpu_percent_list = []
    memory_percent_list = []
    timestamps = []
    process_name = process.name()

    def update(frame):
        cpu_percent = process.cpu_percent(interval=1)
        memory_percent = process.memory_percent()

        cpu_percent_list.append(cpu_percent)
        memory_percent_list.append(memory_percent)
        timestamps.append(time.time() - start_time)

        plt.cla()
        # plt.subplot(2, 1, 1)
        plt.plot(timestamps, cpu_percent_list, label='CPU Usage')
        plt.axhline(y=sum(cpu_percent_list) / len(cpu_percent_list), color='r', linestyle='--', label='CPU Average')
        # plt.xlabel('Time(s)')
        # plt.ylabel('Usage(%)')
        # plt.title(f'{process_name} CPU Usage Over Time')
        # plt.legend()

        # plt.subplot(2, 1, 2)
        plt.plot(timestamps, memory_percent_list, label='Memory Usage')
        plt.axhline(y=sum(memory_percent_list) / len(memory_percent_list), color='g', linestyle='--', label='Memory Average')
        plt.xlabel('Time(s)')
        plt.ylabel('Usage(%)')
        plt.title(f'{process_name} CPU and Memory Usage Over Time')
        plt.legend()

        # plt.tight_layout()

    start_time = time.time()
    fig, ax = plt.subplots()
    ani = FuncAnimation(fig, update, interval=1000)

    plt.show()


if __name__ == "__main__":
    # 解析命令行参数
    parser = argparse.ArgumentParser()
    parser.add_argument('pid', type=int, help='Enter the process id to track')
    args = parser.parse_args()

    # 替换为你要监控的进程ID
    target_process_id = args.pid
    # 监控持续时间（秒）
    # monitoring_duration = 60
    monitor_cpu_memory(target_process_id)
